package sample.methodsearch;

import org.apache.commons.lang3.StringUtils;
import org.objectweb.asm.*;
import org.objectweb.asm.util.Printer;

import java.io.IOException;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import static org.objectweb.asm.Opcodes.ACC_ABSTRACT;
import static org.objectweb.asm.Opcodes.ACC_NATIVE;

public class MethodFindRefDemo {

    public static void main(String[] args) throws IOException {
        invoke(HelloWorld.class, "test", "(III)V");
    }

    private static void invoke(Class<?> clazz, String methodName, String methodDesc) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodFindRefVisitor methodFindInvokeVisitor = new MethodFindRefVisitor(Opcodes.ASM9, classWriter, Type.getInternalName(clazz), methodName, methodDesc);
        classReader.accept(methodFindInvokeVisitor, ClassReader.EXPAND_FRAMES);
    }

    public static class MethodFindRefVisitor extends ClassVisitor {
        protected static final Set<String> METHOD_FIND_REFS = ConcurrentHashMap.newKeySet();
        protected final String methodOwner;
        protected final String methodName;
        protected final String methodDesc;
        private String owner;

        public MethodFindRefVisitor(int api, ClassVisitor classVisitor, String methodOwner, String methodName, String methodDesc) {
            super(api, classVisitor);
            this.methodOwner = methodOwner;
            this.methodName = methodName;
            this.methodDesc = methodDesc;
        }

        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            super.visit(version, access, name, signature, superName, interfaces);
            this.owner = name;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            boolean isAbstractMethod = (access & ACC_ABSTRACT) != 0;
            boolean isNativeMethod = (access & ACC_NATIVE) != 0;
            if (!isAbstractMethod && !isNativeMethod) {
                return new MethodFindRefAdaptor(api, null, owner, name, descriptor, methodOwner, methodName, methodDesc);
            }
            return null;
        }

        @Override
        public void visitEnd() {
            System.out.println(METHOD_FIND_REFS);
            super.visitEnd();
        }

        private static class MethodFindRefAdaptor extends MethodVisitor {

            private String currentOwner;
            private String currentName;
            private String currentDescriptor;

            private String methodOwner;
            private String methodName;
            private String methodDesc;

            public MethodFindRefAdaptor(int api, MethodVisitor methodVisitor,
                                        String currentOwner, String currentName, String currentDescriptor,
                                        String methodOwner, String methodName, String methodDesc) {
                super(api, methodVisitor);
                this.currentOwner = currentOwner;
                this.currentName = currentName;
                this.currentDescriptor = currentDescriptor;
                this.methodOwner = methodOwner;
                this.methodName = methodName;
                this.methodDesc = methodDesc;
            }

            @Override
            public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
                if (StringUtils.equals(owner, methodOwner) && StringUtils.equals(name, methodName) && StringUtils.equals(descriptor, methodDesc)) {
                    String info = String.format("%s %s %s.%s", Printer.OPCODES[opcode], currentOwner, currentName, currentDescriptor);
                    METHOD_FIND_REFS.add(info);
                }
                super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
            }
        }
    }
}
